"""
一些综述的样子

A Survey on LoRA of Large Language Models https://arxiv.org/abs/2407.11046

A Survey on Data Selection for LLM Instruction Tuning https://arxiv.org/abs/2402.05123

Parameter-Efficient Fine-Tuning for Large Models: A Comprehensive Survey https://arxiv.org/abs/2403.14608

"""

import os

from pydantic import BaseModel, Field


class PaperChunk(BaseModel):
    id: int = Field(..., description="论文片段的唯一标识符")
    paper_id: str = Field(..., description="论文的唯一标识符")
    paper_title: str = Field(..., description="论文的标题")
    chunk_id: int = Field(
        ...,
        description="该论文片段在论文中的位置id， 从0开始，可以在回溯的时候使用，比如某一段是根据某个论文的某个chunk_id生成的，那么就可以回溯到原始的论文",
    )
    chunk_text: str = Field(
        ...,
        description="论文片段的文本内容，可能是中文也可能是英文, 中文内容特别少，主要还是英文论文",
    )
    original_filename: str = Field(
        ...,
        description="会议名称,Conf_Paper_Meta_Data_NeurIPS_2022_Neural_Information_Processing_Systems_with_whole_text.db， 意思是NeurIPS的2022年的论文",
    )
    distance: float = Field(
        ...,
        description="api给到的相似度得分，这个得分问题很大，没啥用",
        default_factory=float,
    )

    score: float = Field(..., description="论文片段的相似度分数", default_factory=float)
    clean_text: str = Field(
        ...,
        description="通过大模型压缩过的论文片段， 提取到和topic相关的信息",
        default_factory=str,
    )


    @property
    def text(self):
        return f"\nyear: {self.year} \n title: {self.paper_title} \n text: {self.chunk_text}\n"

    @property
    def reference(self):
        """自己的引用，先只搞标题吧

        标题+所在会议/期刊名称+年份+chunk序号
        [1] Attention is all you need, NeuralIPS, 2022, chunk 0 [2] Attention is all you need, NeuralIPS, 2022, chunk 1
        """
        return f"{self.paper_title}, {self.meeting_name}, {self.year}, chunk {self.chunk_id}. "

    @property
    def year(self):
        import re

        res = re.search(r"\d{4}", self.original_filename)
        if res:
            return res.group(0)
        else:
            return ""

    @property
    def meeting_name(self):
        """
        会议名称或者期刊名称 Conf开头是会议， Journal开头是期刊
        """
        meeting_name = self.original_filename.replace("_with_whole_text.db", "")
        meeting_name = meeting_name.replace("_Paper_Meta_Data", "")
        return meeting_name

    def __hash__(self):
        return hash(self.id)

    def __eq__(self, other):
        if not isinstance(other, PaperChunk):
            return False
        return self.id == other.id

    def __str__(self):
        return self.text

class Paper(BaseModel):
    """
    论文模型：

    1. 对一个**技术概念**的调研综述；
    2. 对一个**方向研究现状**的综述；
    3. 对多个**方法的对比分析**综述；
    4. 对一个**技术方法研究脉络**的综述；
    """

    paper_id: str = Field(..., description="综述的论文ID", default_factory=str)
    paper_title: str = Field(..., description="综述的论文标题", default_factory=str)

    introduction: str = Field(
        ...,
        description="""综述的介绍：
        1. 每一个 section 会去讲什么， 及综述的具体结构
        3. 会绘制综述的结构图

                              """,
        default_factory=str,
    )

    background: str = Field(
        ...,
        description="""综述的背景：
        1. 综述的背景， 及综述的背景
        2. 综述的背景， 及综述的背景
        3. 综述的背景， 及综述的背景
        """,
        default_factory=str,
    )

    conclusion_and_future: str = Field(
        ...,
        description="结论与未来方向, 挑战与方向，一般都以这个结尾",
        default_factory=str,
    )
    references: list = Field(..., description="综述的参考文献", default_factory=list)
    chunks: list[PaperChunk] = Field(
        ..., description="综述的论文片段", default_factory=list
    )

    @property
    def text(self):
        return "\n\n-------------------------------------------------------------------------------------\n\n".join(
            [chunk.chunk_text for chunk in self.chunks]
        )

    @property
    def abstract(self):
        """先默认论文的第一个chunk就是摘要

        Returns:
            _type_: _description_
        """
        return self.chunks[0].chunk_text

    @property
    def original_filename(self):
        return self.chunks[0].original_filename

    @property
    def paper_type(self):
        """
        论文类型：
        Conf: 会议论文
        Journal: 期刊论文
        """
        return self.original_filename.split("_")[0]

    @property
    def title(self):
        return self.chunks[0].paper_title

    def save_to_local(self):
        second_dir = self.original_filename.replace("_with_whole_text.db", "")
        second_dir = second_dir.replace("_Paper_Meta_Data", "")

        save_path = f"data/{second_dir}/{self.title}.md"
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        with open(save_path, "w", encoding="utf-8") as f:
            f.write(self.text)


class Reference(BaseModel):
    """
    参考文献模型
    """

    title: str = Field(..., description="参考文献的标题", default_factory=str)
    author: list = Field(..., description="作者", default_factory=list)
    venue: str = Field(..., description="期刊或会议名称", default_factory=str)
    pages: str = Field(
        ...,
        description="发表信息：卷号和期号，页码，DOI或arXiv编号 ",
        default_factory=str,
    )
    year: str = Field(..., description="发表年份", default_factory=str)


class Subsection(BaseModel):
    title: str = Field(..., description="二级目录标题", default_factory=str)
    description: str = Field(..., description="二级目录描述", default_factory=str)
    content: str = Field(..., description="二级目录内容", default_factory=str)
    paper_chunks: list[PaperChunk] = Field(
        ...,
        description="本章节相关的论文片段",
        default_factory=list,
    )


class Section(BaseModel):
    """
    综述的章节
    """

    title: str = Field(..., description="章节标题", default_factory=str)
    description: str = Field(..., description="章节描述", default_factory=str)
    content: str = Field(..., description="章节内容", default_factory=str)
    subsections: list[Subsection] = Field(
        ...,
        description="二级目录列表，每个二级目录包含标题和描述",
        default_factory=list,
    )

    paper_chunks: list[PaperChunk] = Field(
        ...,
        description="本章节相关的论文片段",
        default_factory=list,
    )


class Survey(BaseModel):
    """
    综述模型：

    1. 对一个**技术概念**的调研综述；
    2. 对一个**方向研究现状**的综述;
    3. 对多个**方法的对比分析**综述;
    4. 对一个**技术方法研究脉络**的综述;
    """

    title: str = Field(..., description="综述的标题", default_factory=str)
    sections: list[Section] = Field(..., description="综述的章节", default_factory=list)
    chunks: list[PaperChunk] = Field(
        ..., description="综述的章节", default_factory=list
    )

    @property
    def outline(self):
        res = ""
        res += f"# {self.title}\n"
        i = 1
        for section in self.sections:
            res += f"## {i}. {section.title}\n"
            j = 1
            for subsection in section.subsections:
                res += f"### {i}.{j}. {subsection.title}\n"
                j += 1
            i += 1
        return res

    @property
    def text(self):
        res = ""
        res += f"\n# {self.title}\n"
        i = 1
        for section in self.sections:
            res += f"\n## {i}. {section.title}\n\n"
            # res += f"{section.content}\n"
            j = 1
            for subsection in section.subsections:
                res += f"\n### {i}.{j}. {subsection.title}\n"
                res += f"{subsection.content}\n"
                j += 1
            i += 1
        # 添加参考文献
        res += self.references
        return res

    @property
    def references(self):
        res = "\n\n ## References\n"
        for i, chunk in enumerate(self.chunks):
            res += f"\n[{i + 1}]  {chunk.reference}\n\n"
        return res
